import pinecone
from pinecone import ServerlessSpec
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from collections import Counter
import time

# -------------------------- 1. 初始化 Pinecone --------------------------
pinecone_api_key = "pcsk_4Nwk78_Q2Pt2WuX6bWXkBUkmxUqGGsMeqDGEjqmiJbjCyTMVUXWVa2f7kooNFbUnJW1W8R"
pinecone_env = "us-east-1"
pc = pinecone.Pinecone(api_key=pinecone_api_key)

# -------------------------- 2. 连接或创建索引 --------------------------
index_name = "mnist-index1"
existing_indexes = pc.list_indexes().names()

if index_name not in existing_indexes:
    print(f"创建索引 {index_name}...")
    pc.create_index(
        name=index_name,
        dimension=64,
        metric="euclidean",
        spec=ServerlessSpec(cloud="aws", region=pinecone_env)
    )
    # 等待索引就绪
    print("等待索引初始化...")
    while True:
        index_desc = pc.describe_index(index_name)
        if index_desc.status.state == "Ready":
            print("索引已就绪！")
            break
        time.sleep(3)
else:
    print(f"索引 {index_name} 已存在，直接连接...")

index = pc.Index(index_name)

# -------------------------- 3. 加载并导入 MNIST 数据 --------------------------
print("\n开始导入 MNIST 数据...")
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

vectors = []
for i in range(len(X)):
    vector_id = str(i)
    vector_values = X[i].tolist()
    metadata = {"label": int(y[i])}
    vectors.append((vector_id, vector_values, metadata))

batch_size = 1000
for i in range(0, len(vectors), batch_size):
    batch = vectors[i:i + batch_size]
    index.upsert(vectors=batch)
    print(f"已导入第 {i//batch_size + 1} 批数据")
print("MNIST 数据导入完成！")

# -------------------------- 4. 自定义数字 3 的 KNN 搜索与预测 --------------------------
digit_3 = np.array(
    [[0, 0, 255, 255, 255, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 0, 255, 255, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 0, 0, 0, 255, 0, 0],
     [0, 0, 255, 255, 255, 255, 0, 0]]
)

digit_3_flatten = (digit_3 / 255.0) * 16
query_data = digit_3_flatten.ravel().tolist()

results = index.query(
    vector=query_data,
    top_k=11,
    include_metadata=True
)

labels = [match.metadata['label'] for match in results.matches]

# 错误处理：检查labels是否为空
if not labels:
    print("未获取到匹配结果，请检查索引数据或查询向量是否正常！")
else:
    for match, label in zip(results.matches, labels):
        print(f"id: {match.id}, distance: {match.score}, label: {label}")

    final_prediction = Counter(labels).most_common(1)[0][0]

    plt.imshow(digit_3, cmap='gray')
    plt.title(f"Predicted digit: {final_prediction}", size=15)
    plt.axis('off')
    plt.show()